import torch


class ValidatorManger(object):
    def __init__(self, model, test_loader, device, criterion):
        self._model = model
        self._test_loader = test_loader
        self.device = device

        self._criterion = criterion

    def validate(self, text):
        self._model.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in self._test_loader:
                data, target = data.to(self.device), target.to(self.device)
                if text:
                    output, state = self._model(data)
                    loss = self._criterion(output.squeeze(), target.squeeze())
                    test_loss += loss.sum() / loss.numel()
                    diff = torch.abs(output.squeeze() - target.squeeze())  # 取绝对值
                    # 判断哪些差异小于 0.1
                    correct_ = (diff < 0.2).sum().item()
                    # 准确率
                    correct += correct_ / target.numel()
                else:
                    output = self._model(data)
                    test_loss += self._criterion(output, target).item()
                    pred = output.argmax(dim=1, keepdim=True)
                    correct += pred.eq(target.view_as(pred)).sum().item()

        test_loss /= len(self._test_loader.dataset)
        accuracy = 100. * correct / len(self._test_loader.dataset)
        print(
            f"Test set: Average loss: {test_loss}, Accuracy: {correct}/{len(self._test_loader.dataset)} ({accuracy})")
